package tools

import (
	"FscanProjectTools/common"
	"fmt"
	"sync"
	"time"
)

// 变量
var (
	WorkerThreads int = 100
	ChanNum       int = 160
)

type Addr struct {
	ip   string
	port int
}

type Portscaner struct {
	HostsList []string
	PortsList []int
	IpPorts   map[string][]int
	Timeout   int64
}

func NewPortscaner(hosts []string, ports string, timeout int64) (portscaner Portscaner) {
	portscaner.HostsList = hosts
	portscaner.PortsList = common.ParsePort(ports)
	portscaner.IpPorts = make(map[string][]int, len(hosts))
	portscaner.Timeout = timeout
	return portscaner
}

func (psr *Portscaner) Run() map[string][]int {
	Addrs := make(chan Addr, ChanNum)
	results := make(chan Addr, ChanNum)
	var wg sync.WaitGroup

	// 消费者 接收结果
	go func() {
		for result := range results {
			//psr.AliveAddress = append(psr.AliveAddress, result)
			if _, ok := psr.IpPorts[result.ip]; !ok {
				psr.IpPorts[result.ip] = []int{}
			}
			psr.IpPorts[result.ip] = append(psr.IpPorts[result.ip], result.port)
			wg.Done()
		}
	}()
	// 多线程扫描
	for i := 0; i < WorkerThreads; i++ {
		go func() {
			for addr := range Addrs {
				PortConnect(addr, results, psr.Timeout, &wg)
				wg.Done()
			}
		}()
	}

	// 添加扫描目标
	for _, port := range psr.PortsList {
		for _, host := range psr.HostsList {
			wg.Add(1)
			Addrs <- Addr{host, port}
		}
	}
	wg.Wait()
	close(Addrs)
	close(results)
	return psr.IpPorts
}

func PortConnect(addr Addr, respondingHosts chan<- Addr, adjustedTimeout int64, wg *sync.WaitGroup) {
	host, port := addr.ip, addr.port
	conn, err := common.WrapperTcpWithTimeout("tcp4", fmt.Sprintf("%s:%v", host, port), time.Duration(adjustedTimeout)*time.Second)
	if err == nil {
		defer conn.Close()
		wg.Add(1)
		respondingHosts <- Addr{host, port}
	}
}
